import unittest
from unittest.mock import patch, call
from Utils.RedisTokenPool import RedisTokenPool


class TestRedisTokenPool(unittest.TestCase):

    def setUp(self):
        self.tokens_key = "test_tokens"
        self.tokens = ["token1", "token2", "token3"]

    @patch('redis.StrictRedis')
    def test_initialization(self, mock_redis):
        # 模拟 Redis 实例
        mock_redis_instance = mock_redis.return_value
        mock_redis_instance.exists.return_value = False

        # 创建 RedisTokenPool 实例
        RedisTokenPool(self.tokens_key, self.tokens)

        # 验证是否向 Redis 中添加了令牌
        expected_calls = [call(self.tokens_key, token) for token in reversed(self.tokens)]
        mock_redis_instance.lpush.assert_has_calls(expected_calls, any_order=True)

    @patch('redis.StrictRedis')
    def test_get_token(self, mock_redis):
        mock_redis_instance = mock_redis.return_value
        mock_redis_instance.brpoplpush.return_value = "token1"
        token_pool = RedisTokenPool(self.tokens_key, self.tokens)

        with token_pool.get_token() as token:
            self.assertEqual(token, "token1")

    @patch('redis.StrictRedis')
    def test_get_token_timeout(self, mock_redis):
        # 模拟 Redis 实例
        mock_redis_instance = mock_redis.return_value
        mock_redis_instance.brpoplpush.side_effect = [None] * 5  # 模拟 5 次尝试都超时

        token_pool = RedisTokenPool(self.tokens_key, self.tokens)

        # 测试超时异常
        with self.assertRaises(TimeoutError):
            with token_pool.get_token():
                pass


if __name__ == '__main__':
    unittest.main()
